import threading
import time
from collections import deque
from typing import Dict, Optional
import logging


class PerformanceMonitor(threading.Thread):
    """性能监控线程类"""
    
    def __init__(self, update_interval: float = 2.0, logger: Optional[logging.Logger] = None):
        """
        初始化性能监控器
        
        Args:
            update_interval: 统计更新间隔（秒）
            logger: 日志记录器，如果为None则使用print
        """
        super().__init__(name="PerformanceMonitor", daemon=True)
        
        self.update_interval = update_interval
        self.logger = logger
        self.shutdown_flag = False
        
        # 统计数据锁
        self.stats_lock = threading.Lock()
        
        # 各阶段计数器和时间记录
        self.read_count = 0
        self.detection_count = 0
        self.position_count = 0
        self.visualization_count = 0
        
        # 时间记录队列（保存最近30个时间点）
        self.read_times = deque(maxlen=30)
        self.detection_times = deque(maxlen=30)
        self.position_times = deque(maxlen=30)
        self.visualization_times = deque(maxlen=30)
        
        # 启动时间
        self.start_time = time.time()
        self.last_stats_update = time.time()
        
        # 最近的FPS统计
        self.latest_fps = {
            'read': 0.0,
            'detection': 0.0,
            'position': 0.0,
            'visualization': 0.0
        }
    
    def update_stats(self, stage: str, success: bool = True) -> None:
        """
        更新指定阶段的统计信息
        
        Args:
            stage: 阶段名称 ('read', 'detection', 'position', 'visualization')
            success: 是否成功（对于detection和position有意义）
        """
        current_time = time.time()
        
        with self.stats_lock:
            if stage == 'read':
                self.read_count += 1
                self.read_times.append(current_time)
            elif stage == 'detection' and success:
                self.detection_count += 1
                self.detection_times.append(current_time)
            elif stage == 'position' and success:
                self.position_count += 1
                self.position_times.append(current_time)
            elif stage == 'visualization':
                self.visualization_count += 1
                self.visualization_times.append(current_time)
    
    def calculate_fps(self, time_deque: deque) -> float:
        """计算FPS"""
        if len(time_deque) < 2:
            return 0.0
        time_span = time_deque[-1] - time_deque[0]
        if time_span > 0:
            return (len(time_deque) - 1) / time_span
        return 0.0
    
    def get_current_stats(self) -> Dict:
        """获取当前统计信息"""
        with self.stats_lock:
            read_fps = self.calculate_fps(self.read_times)
            detection_fps = self.calculate_fps(self.detection_times)
            position_fps = self.calculate_fps(self.position_times)
            visualization_fps = self.calculate_fps(self.visualization_times)
            
            # 计算成功率
            detection_rate = (self.detection_count / max(self.read_count, 1)) * 100
            position_rate = (self.position_count / max(self.detection_count, 1)) * 100
            visualization_rate = (self.visualization_count / max(self.read_count, 1)) * 100
            
            # 更新最新FPS
            self.latest_fps = {
                'read': read_fps,
                'detection': detection_fps,
                'position': position_fps,
                'visualization': visualization_fps
            }
            
            return {
                'fps': self.latest_fps.copy(),
                'counts': {
                    'read': self.read_count,
                    'detection': self.detection_count,
                    'position': self.position_count,
                    'visualization': self.visualization_count
                },
                'rates': {
                    'detection': detection_rate,
                    'position': position_rate,
                    'visualization': visualization_rate
                },
                'total_time': time.time() - self.start_time
            }
    
    def find_bottleneck(self, fps_dict: Dict[str, float]) -> tuple:
        """找出性能瓶颈"""
        bottleneck = min(fps_dict, key=fps_dict.get)
        bottleneck_fps = fps_dict[bottleneck]
        
        # 生成优化建议
        suggestions = {
            'read': "优化图像读取，考虑并行读取或更快的存储",
            'detection': "增加seg_workers和pose_workers数量，或优化模型",
            'position': "优化compute_position算法或增加处理线程",
            'visualization': "简化可视化效果或降低显示分辨率"
        }
        
        return bottleneck, bottleneck_fps, suggestions.get(bottleneck, "未知优化建议")
    
    def log_stats(self, stats: Dict) -> None:
        """记录统计信息"""
        fps = stats['fps']
        counts = stats['counts']
        rates = stats['rates']
        
        # 找出瓶颈
        bottleneck, bottleneck_fps, suggestion = self.find_bottleneck(fps)
        
        # 格式化输出
        output_lines = [
            "=== 性能统计 ===",
            f"读图片:   {fps['read']:.2f} FPS | 总数: {counts['read']}",
            f"检测:     {fps['detection']:.2f} FPS | 总数: {counts['detection']} | 成功率: {rates['detection']:.1f}%",
            f"位置计算: {fps['position']:.2f} FPS | 总数: {counts['position']} | 成功率: {rates['position']:.1f}%",
            f"可视化:   {fps['visualization']:.2f} FPS | 总数: {counts['visualization']} | 处理率: {rates['visualization']:.1f}%",
            f"性能瓶颈: {bottleneck} ({bottleneck_fps:.2f} FPS)",
            "=" * 50
        ]
        
        for line in output_lines:
            if self.logger:
                self.logger.info(line)
            else:
                print(line)
    
    def run(self) -> None:
        """监控线程主循环"""
        while not self.shutdown_flag:
            try:
                time.sleep(self.update_interval)
                
                # 获取并输出统计信息
                stats = self.get_current_stats()
                self.log_stats(stats)
                
            except Exception as e:
                if self.logger:
                    self.logger.error(f"性能监控错误: {e}")
                else:
                    print(f"性能监控错误: {e}")
    
    def get_latest_fps(self, stage: str) -> float:
        """获取指定阶段的最新FPS"""
        return self.latest_fps.get(stage, 0.0)
    
    def shutdown(self) -> None:
        """关闭监控线程"""
        self.shutdown_flag = True
    
    def get_final_report(self) -> str:
        """获取最终性能报告"""
        stats = self.get_current_stats()
        fps = stats['fps']
        counts = stats['counts']
        rates = stats['rates']
        total_time = stats['total_time']
        
        bottleneck, bottleneck_fps, suggestion = self.find_bottleneck(fps)
        
        report_lines = [
            "=" * 60,
            "最终性能统计",
            "=" * 60,
            f"总运行时间: {total_time:.2f}s",
            "",
            f"读图片:   {fps['read']:.2f} FPS | 总数: {counts['read']}",
            f"检测:     {fps['detection']:.2f} FPS | 总数: {counts['detection']} | 成功率: {rates['detection']:.1f}%",
            f"位置计算: {fps['position']:.2f} FPS | 总数: {counts['position']} | 成功率: {rates['position']:.1f}%",
            f"可视化:   {fps['visualization']:.2f} FPS | 总数: {counts['visualization']}",
            "",
            f"性能瓶颈: {bottleneck} ({bottleneck_fps:.2f} FPS)",
            f"优化建议: {suggestion}",
            "=" * 60
        ]
        
        return "\n".join(report_lines)